package com.rt.gptnext.service;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.TypeReference;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.rt.gptnext.common.constant.HttpConstants;
import com.rt.gptnext.common.utils.http.HttpResponse;
import com.rt.gptnext.common.utils.http.HttpUtil;
import com.rt.gptnext.common.utils.io.MultiOutputStream;
import com.rt.gptnext.common.utils.json.JsonUtil;
import com.rt.gptnext.data.entity.OpenaiChat;
import com.rt.gptnext.data.entity.OpenaiCompletion;
import com.rt.gptnext.data.enums.OpenaiCompletionRoleEnum;
import com.rt.gptnext.data.mapper.OpenaiCompletionMapper;
import com.rt.gptnext.properties.OpenaiProperties;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import java.io.ByteArrayOutputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author gaoyiyang
 */
@Service
@AllArgsConstructor
@Slf4j
public class OpenaiCompletionService extends ServiceImpl<OpenaiCompletionMapper, OpenaiCompletion> {

    private final OpenaiProperties openaiProperties;

    private final Set<String> removeQaIdSet = new CopyOnWriteArraySet<>();

    public void remove(String qaId, Long id) {
        if (id != null) {
            removeById(id);
            return;
        }
        if (StrUtil.isNotBlank(qaId)) {
            if (!lambdaUpdate()
                    .eq(OpenaiCompletion::getQaId, qaId)
                    .remove()) {
                // 删除失败加入临时删除set，用以阻止入库
                removeQaIdSet.add(qaId);
            }
        }
    }

    public List<OpenaiCompletion> completions(Long chatId, Integer version) {
        return lambdaQuery()
                .eq(OpenaiCompletion::getChatId, chatId)
                .eq(OpenaiCompletion::getChatVersion, version)
                .ne(OpenaiCompletion::getRole, OpenaiCompletionRoleEnum.SYSTEM)
                .list();
    }

    public String completion(OpenaiChat chat, String content, String questionId, String answerId) {
        return sendCompletion(questionId, answerId, chat, content, false, null);
    }

    public void streamCompletion(OpenaiChat chat, String content, OutputStream out, String questionId, String answerId) {
        sendCompletion(questionId, answerId, chat, content, true, out);
    }

    @Transactional(rollbackFor = Throwable.class)
    public String reset(OpenaiChat chat, Integer copyVersion) {
        var list = lambdaQuery()
                .eq(OpenaiCompletion::getChatId, chat.getId())
                .eq(OpenaiCompletion::getChatVersion, copyVersion)
                .list();
        if (CollUtil.isEmpty(list)) {
            return null;
        }
        // 最后一个问题下标
        var lastQuestionIdx = CollUtil.lastIndexOf(list, c -> c.getRole() == OpenaiCompletionRoleEnum.USER);
        Long lastQuestionId = null;
        for (int i = 0; i < lastQuestionIdx; i++) {
            var completion = list.get(i);
            if (i == list.size() - 1 && completion.getRole() == OpenaiCompletionRoleEnum.ASSISTANT) {
                continue;
            }
            OpenaiCompletion newCompletion = new OpenaiCompletion();
            newCompletion.setChatVersion(chat.getVersion());
            newCompletion.setChatId(chat.getId());
            newCompletion.setRole(completion.getRole());
            newCompletion.setContent(completion.getContent());
            save(newCompletion);
            if (newCompletion.getRole() == OpenaiCompletionRoleEnum.USER) {
                lastQuestionId = newCompletion.getId();
            }
            if (newCompletion.getRole() == OpenaiCompletionRoleEnum.ASSISTANT) {
                newCompletion.setQuestionId(lastQuestionId);
            }
        }
        var lastQuestion = CollUtil.get(list, lastQuestionIdx);
        if (lastQuestion == null) {
            return null;
        }
        return lastQuestion.getContent();
    }

    private String sendCompletion(String questionId, String answerId, OpenaiChat chat, String content, boolean stream, OutputStream out) {
        var completions = lambdaQuery()
                .eq(OpenaiCompletion::getChatId, chat.getId())
                .eq(OpenaiCompletion::getChatVersion, chat.getVersion())
                .list();
        // 如果是第一次问答，拼接系统prompt到最前端
        if (CollUtil.isEmpty(completions) && StrUtil.isNotBlank(chat.getSystemPrompt())) {
            var systemCompletion = new OpenaiCompletion();
            systemCompletion.setRole(OpenaiCompletionRoleEnum.SYSTEM);
            systemCompletion.setContent(chat.getSystemPrompt());
            systemCompletion.setChatId(chat.getId());
            systemCompletion.setChatVersion(chat.getVersion());
            save(systemCompletion);
            completions.add(systemCompletion);
        }
        // 拼接最新的问题
        var question = new OpenaiCompletion();
        question.setQaId(questionId);
        question.setRole(OpenaiCompletionRoleEnum.USER);
        question.setContent(content);
        question.setChatId(chat.getId());
        question.setChatVersion(chat.getVersion());
        save(question);
        completions.add(question);

        // 画图
        if (isDraw(content)) {
            return draw(answerId, question, chat, stream, out);
        }

        // 将completions转为请求体
        List<JSONObject> messages = new ArrayList<>();
        for (var completion : completions) {
            if (completion.isDraw()) {
                continue;
            }
            messages.add(JSONObject.of("role", completion.getRole().getValue(), "content", completion.getContent()));
        }
        var body = JSONObject.of("model", "gpt-3.5-turbo")
                .fluentPut("temperature", chat.getTemperature())
                .fluentPut("messages", messages)
                .fluentPut("stream", stream);

        // 组装请求对象
        var request = HttpUtil.request(openaiProperties.getUrl() + "/v1/chat/completions")
                .post()
                .authBearer(pollKey())
                .header("Content-Type", HttpConstants.MimeTypes.JSON)
                .header("Accept", HttpConstants.MimeTypes.JSON)
                .body(body);

        // 回答
        String answerContent;

        // 发送请求
        if (stream) {
            // 流式请求
            // 组装响应流，以保存响应数据
            var bos = new ByteArrayOutputStream();
            var mos = new MultiOutputStream(bos, out);
            request.streamRead(msg -> {
                try {
                    if (StrUtil.isBlank(msg)) {
                        return true;
                    }
                    if (!StrUtil.startWith(msg, "data:")) {
                        IoUtil.write(mos, StandardCharsets.UTF_8, false, msg);
                        return true;
                    }
                    var data = JSONObject.parse(msg.replace("data:", ""));
                    var choices = JsonUtil.getByJsonPath(data, "$.choices", new TypeReference<List<JSONObject>>() {
                    });
                    for (JSONObject choice : choices) {
                        if ("stop".equals(choice.getString("finish_reason"))) {
                            return false;
                        }
                        var str = JsonUtil.getByJsonPath(choice, "$.delta.content", String.class);
                        if (str != null) {
                            IoUtil.write(mos, StandardCharsets.UTF_8, false, str);
                        }
                    }
                    return true;
                } catch (Exception e) {
                    log.error("openai流式对话报错", e);
                    return false;
                }
            }, true).blocking();
            IoUtil.close(out);
            answerContent = bos.toString(StandardCharsets.UTF_8);
        } else {
            // 非流式请求
            var response = request.validator(resp -> StrUtil.isNotBlank(resp.getByJsonPath("$.choices[0].message.content", String.class)))
                    .call();
            // 回答
            if (!response.isValidateSuccess()) {
                answerContent = "请求失败,未知异常";
            } else {
                answerContent = response.getByJsonPath("$.choices[0].message.content", String.class);
            }
        }

        // 保存回答
        var answer = new OpenaiCompletion();
        answer.setContent(answerContent);
        // 问题ID保存
        answer.setQuestionId(question.getId());
        answer.setQaId(answerId);
        answer.setChatId(chat.getId());
        answer.setRole(OpenaiCompletionRoleEnum.ASSISTANT);
        answer.setChatVersion(chat.getVersion());
        if (!removeQaIdSet.remove(answer.getQaId())) {
            save(answer);
        }
        return answerContent;
    }

    private String draw(String answerId, OpenaiCompletion question, OpenaiChat chat, boolean stream, OutputStream out) {
        OutputStream bos = new ByteArrayOutputStream();
        OutputStream mos = stream ? new MultiOutputStream(bos, out) : bos;
        IoUtil.writeUtf8(mos, false, "正在生成图片，请耐心等待...");
        var resp = HttpUtil.request(openaiProperties.getUrl() + "/v1/images/generations")
                .post()
                .authBearer(pollKey())
                .body(JSONObject.of("prompt", StrUtil.removeAny(question.getContent(), "画图：", "画图:", "draw:", "draw："), "n", 1, "response_format", "url"))
                .call();
        var result = resp.getByJsonPath("$.data[0].url", String.class);
        IoUtil.writeUtf8(mos, false, "\n生成完成:\n");
        IoUtil.writeUtf8(mos, false, "![image](" + result + ")");
        String answerContent = bos.toString();
        // 保存回答
        var answer = new OpenaiCompletion();
        answer.setContent(answerContent);
        // 问题ID保存
        answer.setQuestionId(question.getId());
        answer.setQaId(answerId);
        answer.setChatId(chat.getId());
        answer.setRole(OpenaiCompletionRoleEnum.ASSISTANT);
        answer.setChatVersion(chat.getVersion());
        answer.setDraw(true);
        if (!removeQaIdSet.remove(answer.getQaId())) {
            save(answer);
        }
        // 更新问题
        question.setDraw(true);
        updateById(question);
        return answerContent;
    }


    private final AtomicInteger loop = new AtomicInteger();

    private String pollKey() {
        if (CollUtil.isEmpty(openaiProperties.getKey())) {
            return null;
        }
        var idx = loop.getAndIncrement() % openaiProperties.getKey().size();
        return openaiProperties.getKey().get(idx);
    }

    private boolean isDraw(String content) {
        if (StrUtil.isBlank(content)) {
            return false;
        }
        var str = StrUtil.trim(content);
        return StrUtil.startWithAnyIgnoreCase(str, "画图：", "画图:", "draw:", "draw：");
    }

}
